import yaml
import pandas as pd
from itertools import product
from time import time

from tjd_config.ini_data import ini_y_data
from tjd_config.ini_data import ini_x_data


def load_config():
    with open("config.yaml", 'r', encoding='utf-8') as f:
        data = yaml.load(f.read(), yaml.FullLoader)
    return data


class SubEdges(object):
    def __init__(self, edges, vertex=None, *args, **kwargs):
        self.edges = edges
        self.vertex = vertex if vertex else next(iter(self.edges))
        self.subedges = []

    def __call__(self, *args, **kwargs):
        b = time()
        print('Start SubEdges for {}.'.format(self.vertex))
        self.start_recursion()
        e = time()
        print('End SubEdges for {}, generate subedges {}, take time {:.4f}s.'.format(self.vertex, len(self.subedges),
                                                                                     e - b))
        return self.subedges

    def start_recursion(self):
        se = []
        self.recursion([self.vertex], se)
        return self.subedges

    def recursion(self, keys, se):
        key, value = [], []
        for i in keys:
            value.append(product(*self.edges[i].values()))
            [key.append(j) for j in self.edges[i].keys()]
        value = product(*value)

        for i in value:
            i = [k for j in i for k in j]
            se_ = se.copy()
            se_.append(dict(zip(key, i)))

            value_bool = [j not in self.edges for j in i]
            if all(value_bool):

                se_ = {k: v for d in se_ for k, v in d.items()}
                self.subedges.append(se_)
            else:
                idx = [i for i, j in enumerate(value_bool) if not j]
                keys = [list(i)[j] for j in idx]
                self.recursion(keys, list(se_))


def get_edges(y_name):
    cfg = load_config()
    edges = {
        'x': ini_x_data(cfg, y_name),
        'y': ini_y_data(cfg),

        'flow': {
            'slicer1': ['Slicer1'],
            'selector': ['Filter', 'Distance'],
            'slicer2': ['Slicer2'],
            'learner': ['machine_learning', 'deep_learning']
        },

        'engineer': {
            'slicer1': ['Slicer1'],
            'selector': ['Filter', 'Correlation_filter', 'Distance'],
        },

        'trainer': {
            'slicer2': ['Slicer2'],
            'learner': ['machine_learning', 'deep_learning'][:1]
        },

        'Slicer1': {
            'y_name': ['银行'],  # (y_dim3)
            'freq': [cfg['freq']],  # (xy_dim1)
            'end_datetime':
                list(pd.date_range(start=cfg['end_datetime_from'],
                                   end=cfg['end_datetime_to'],
                                   freq=cfg['freq']).strftime('%F')),  # (xy_dim1)
            'periods': list(range(*cfg['periods_args'])),  # (xy_dim1)
            'test_len': [1]  # (xy_dim1)
        },

        'Filter': {
            'score_func': ['f_regression', 'mutual_info_regression'][:],
            'mode': ['k_best', 'fpr', 'fdr', 'fwe'][:],
        },

        'Correlation_filter': {
            'method': ['pearson', 'kendall', 'spearman'][:2]
        },

        'Distance': {
            'metric': ['braycurtis', 'canberra', 'chebyshev', 'cityblock', 'correlation', 'cosine', 'dice', 'euclidean',
                       'hamming', 'jaccard'][:]
        },

        'Slicer2': {
            'n_features': list(cfg['n_features_args']),  # (x_dim3)
            'x_time_step': [1],  # (x_dim2)
            'prediction': [False],  # (y_dim1)
            'y_time_step': [1],  # (y_dim2)
            'valid_pct': cfg['valid_pct_args']  # (xy_dim1)
        },

        'machine_learning': {
            'model': ['LinearRegression', 'Ridge', 'Lasso', 'ElasticNet',
                      'BayesianRegression',
                      'AutomaticRelevanceDetermination',
                      'StochasticGradientDescent',
                      ][:],
        },

        'deep_learning': {
            'model': ['RNN', 'GRU', 'LSTM', 'TCN', 'DARNN', 'TransformerEncoderLayer'][3:5],
            'batch_size': [12],
            'loss_func': ['MSELoss'],
            'optimizer': ['adam'],
            'learning_rate': [0.01],
            'epochs': [400],
            'nhidden_encoder': [64],
            'nhidden_decoder': [64]
        },

        'RNN': {
            'hidden_size': [10],
        }

    }
    return edges


if __name__ == '__main__':
    print(666)
    cfg = load_config()
    edges = get_edges()
